package model

import (
	"fmt"
	"gitee.com/jn-qq/big-model/common"
	"gitee.com/jn-qq/big-model/openai"
	"gitee.com/jn-qq/go-tools/data"
	"maps"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"vector/tools"
)

type RagLabelTask struct {
	content      string
	baseResult   string
	ragResult    string
	polishResult string
	resultData   map[string]any
}

func (r *RagLabelTask) Reset() {
	r.content = ""
	r.ragResult = ""
	r.polishResult = ""
	r.baseResult = ""
	clear(r.resultData)
}

func (r *RagLabelTask) NewTaskData(row map[string]any) error {
	// 复制数据
	r.resultData = maps.Clone(row)
	r.content = row["问题"].(string)
	r.ragResult = row["RAG结果"].(string)
	r.polishResult = row["润色结果"].(string)
	r.baseResult = row["参考结果"].(string)
	return nil
}

func (r *RagLabelTask) Do() {
	if r.ragResult == "" || r.ragResult == "/" {
		r.resultData["RAG标注"] = false
		r.resultData["RAG得分"] = -1
		r.resultData["润色标注"] = false
		return
	}
	// RAG标注
	modelAnswer, err := r.Request(r.CreatPrompt(r.content, r.ragResult, r.baseResult))
	if err != nil {
		r.resultData["备注"] = err.Error()
		return
	}

	answer := data.JsonString(modelAnswer)
	if !answer.Check() {
		r.resultData["备注"] = "解析json失败"
		return
	}
	r.resultData["RAG标注"], _ = answer.Find("is_accurate")
	r.resultData["RAG得分"], _ = answer.Find("score")
	if !r.resultData["RAG标注"].(bool) || r.polishResult == "" || r.polishResult == "/" {
		r.resultData["润色标注"] = false
		return
	}

	// 润色标注
	r.resultData["润色标注"], err = r.Request(fmt.Sprintf(`请判断以下【润色结果】是否与【原数据】的意思一致。请严格遵循以下规则：<ret><ret>1. 核心信息一致性：确保关键事实、数据、名称和结论等内容保持一致。<ret>2. 语义一致性：允许表达方式、句式或词语的变化，但整体语义不应改变。<ret>3. 无信息增减：润色后的文本不应添加或删减原数据中的重要信息。<ret>4. 语气与风格：可适当调整语气或风格（如正式或口语化），但不改变原意。<ret><ret>请特别注意：  <ret>你的输出必须严格遵循指定格式，仅返回一个判断结果，不要包含任何分析、解释或额外内容。<ret><ret>若意思相同，请仅输出：True  <ret>若意思不同，请仅输出：False  <ret>若不确认，请仅输出：Nan  <ret><ret>【原数据】：  <ret>%s  <ret><ret>【润色结果】：  <ret>%s  <ret>`, r.ragResult, r.polishResult))
	if err != nil {
		r.resultData["备注"] = err.Error()
		return
	}

}

func (r *RagLabelTask) Request(prompt string) (string, error) {
	// 获取模型配置
	model, ok := tools.Config.Agent["14B"]
	if !ok {
		panic("未找到对应的模型配置")
	}
	client := openai.NewOpenAiClient(common.Model{
		Url:       model.Address,
		ApiKey:    model.ApiKey,
		ModelType: model.ModeType,
		IsReason:  false,
	}).
		SetStream(model.Stream).
		SetMaxRounds(model.MaxTokens).
		SetTemperature(model.Temperature).
		SetTopK(model.TopK)
	// 标注模型提示词
	response, err := client.Chat(prompt)
	if err != nil {
		return "", err
	}
	return response.GetAnswer(), nil
}

func (r *RagLabelTask) CreatPrompt(question, ragResult, baseResult string) string {
	var prompt string
	if baseResult == "" {
		prompt = `<User> 你是一个智能助手，负责评估检索系统返回的结果质量。<ret><ret>请根据以下信息，评估检索结果在回答用户问题方面的质量，并给出一个 1-5 的评分。<ret><ret>**评分标准：**<ret>- **1分 (非常差)**: 检索结果与问题完全无关，或包含严重错误信息。<ret>- **2分 (较差)**: 检索结果与问题相关性很低，或信息严重不足，无法回答问题。<ret>- **3分 (一般)**: 检索结果部分相关，提供了一些有用信息，但不够完整或准确，未能很好地回答问题。<ret>- **4分 (良好)**: 检索结果基本正确且相关，能较准确地回答大部分问题，细节可能略有不足。<ret>- **5分 (优秀)**: 检索结果完全正确、完整且高度相关，准确无误地回答了用户问题。<ret><ret>**用户问题：**<ret>[USER_QUESTION]<ret><ret>**检索结果：**<ret>[RETRIEVAL_RESULT]<ret><ret>**你的任务：**<ret>请根据上述标准进行评估，然后以严格的 JSON 格式输出结果，包含两个字段：<ret>- "score": 整数，表示 1-5 的评分。<ret>- "is_accurate": 布尔值，如果评分是 4 或 5，则为 true；如果评分是 1, 2 或 3，则为 false。<ret>例如：{"score": 4, "is_accurate": true} 或 {"score": 2, "is_accurate": false}<ret>请只输出这个 JSON 字符串，不要添加任何其他文字或解释。<end><Bot>`
	} else {
		prompt = `<User> 你是一个智能助手，负责评估检索系统返回的结果与标准答案的一致性。<ret><ret>请根据以下信息，评估检索结果在语义上与参考答案的一致程度，并给出一个 1-5 的评分。<ret><ret>**评分标准：**<ret>- **1分 (非常差)**: 检索结果与参考答案的核心信息完全矛盾或毫不相关。<ret>- **2分 (较差)**: 检索结果与参考答案存在明显矛盾，或关键信息严重缺失。<ret>- **3分 (一般)**: 检索结果与参考答案部分一致，但存在一些不一致之处或信息不够完整。<ret>- **4分 (良好)**: 检索结果的核心信息和观点与参考答案基本一致，细节上可能有轻微差异。<ret>- **5分 (优秀)**: 检索结果的语义与参考答案完全一致，表达了相同或高度相似的意思。<ret><ret>**用户问题：**<ret>[USER_QUESTION]<ret><ret>**检索结果：**<ret>[RETRIEVAL_RESULT]<ret><ret>**参考答案：**<ret>[REFERENCE_ANSWER]<ret><ret>**你的任务：**<ret>请根据上述标准进行评估，然后以严格的 JSON 格式输出结果，包含两个字段：<ret>- "score": 整数，表示 1-5 的评分。<ret>- "is_accurate": 布尔值，如果评分是 4 或 5，则为 true；如果评分是 1, 2 或 3，则为 false。<ret>例如：{"score": 5,"is_accurate": true} 或 {"score": 3,"is_accurate": false}<ret>请只输出这个 JSON 字符串，不要添加任何其他文字或解释。<end><Bot>`
	}
	prompt = strings.ReplaceAll(prompt, "[USER_QUESTION]", question)
	prompt = strings.ReplaceAll(prompt, "[RETRIEVAL_RESULT]", ragResult)
	prompt = strings.ReplaceAll(prompt, "[REFERENCE_ANSWER]", baseResult)
	return prompt
}

func (r *RagLabelTask) Result() map[string]any {
	return maps.Clone(r.resultData)
}

type RagTask struct {
	content    string
	question   string
	resultData map[string]any
}

func (r *RagTask) Reset() {
	r.content = ""
	r.question = ""
	clear(r.resultData)
}

func (r *RagTask) NewTaskData(row map[string]any) error {
	// 复制数据
	r.resultData = maps.Clone(row)
	r.content = row["content"].(string)
	r.question = row["问题"].(string)
	return nil
}

func (r *RagTask) Do1() {
	prompt := `<User> 你是一个文档问答专家，请严格根据提供的知识片段中的参考答案进行准确的回复语生成。请按以下步骤处理：<ret># 任务要求<ret>1. 问题分类（必须二选一）：<ret>- [Reply] 用户问题与相似问题语义相关的<ret>- [Reject] 用户问题与相似问题语义完全不相关的<ret>2. 回复标准：<ret>- 使用匹配度最高的参考答案，生成回复语，不允许输出示例内容，不符合条件就返回Reject <ret>- 用户星级和全球通等级没有对应关系,禁止出现四星全球通用户（对应银卡等级）相关话术 <ret>- 绝对禁止编造内容，尤其是数值类<ret>- 回复语不要出现出现其他渠道如：联系官方客服、或者什么移动app之类的<ret>- 答案里有链接地址的，必须对地址进行<a href=' '>请点击这里</a >标签包装，将链接替换这里的url。例如：<a href='xxxx'>请点击这里</a >。<ret>- 输出的内容可以在适当的位置做换行，凸显调理清晰<ret>- 回复答案中不要出现原子主名称/原子名称/知识名称等字段  <ret>#可参考知识如下：%s<ret>用户输入问题：%s<ret># 输出模板（严格遵循）<ret># Type [分类标签]<ret># Response<ret><ret># 输出示例（严格遵循）<ret># Type [Reply]<ret># Response <ret>回复内容 <end><Bot> `
	prompt = fmt.Sprintf(prompt, r.content, r.question)
	client := openai.NewOpenAiClient(common.Model{
		Url:       "http://172.29.100.69:9969/v1/chat/completions",
		ApiKey:    "",
		ModelType: "x",
		IsReason:  false,
	}).
		SetStream(false).
		SetMaxTokens(4096).
		SetTemperature(0.6).
		SetTopK(0.99)
	response, err := client.Chat(prompt)
	if err != nil {
		r.resultData["备注"] = err.Error()
		return
	}
	answer := response.GetAnswer()
	// 正则表达式：匹配 # Response 和 <end> 之间的内容
	// 使用 [\s\S] 或 [\d\D] 来匹配包括换行在内的所有字符
	re := regexp.MustCompile(`Response\s*([\s\S]*?)<end>`)

	// 查找匹配的内容
	matches := re.FindStringSubmatch(answer)

	if len(matches) >= 2 {
		content := strings.TrimSpace(matches[1])
		answer = content
	}
	r.resultData["rag_result"] = answer
}

func (r *RagTask) Do2() {
	prompt := `<User>你是一名专业的文本校对AI助手，名为“积分改名检查官”。您的角色是严格遵循中国移动积分更名为AI豆的规则，对用户提供的文本进行快速、准确的校对。您以严谨、细致著称，只关注文本中“积分”关键字的替换情况，不添加任何额外解释或情感色彩。您的唯一目标是输出“是”或“否”，确保结果客观可靠。<ret># 校对背景<ret>自2025年9月26日起，中国移动积分正式更名为AI豆。需要检查文本中所有“积分”关键字是否已按规则替换为“AI豆”。<ret># 核心规则<ret>1. 必须替换的情况：<ret>- 普通描述中的“积分”（如“您的积分”、“兑换积分”、“积分过期”），必须替换为“AI豆”。<ret>2. 例外情况（保持原样）：<ret>- 【积分商城】作为整体专有名词。<ret>- 其他固定短语或活动名称，如“积分活动”、“信用积分”、“积分计划”（例如：“开通积分计划”中的“积分”不替换）。<ret>- 【中国移动积分正式更名为AI豆】作为给用户的提示里面的积分字眼不用替换。<ret># 操作步骤<ret>1. 输入：接收用户提供的字符串文本。<ret>2. 检查：<ret>- 逐句扫描文本，标识所有“积分”字眼。<ret>- 对每个“积分”，判断是否属于例外列表：<ret>  - 是例外（如“积分商城”）→ 保持原样正确。<ret>  - 非例外（如“您有100积分”）→ 必须已替换为“AI豆”。<ret>3. 输出：<ret>- 如果文本中所有非例外的“积分”均已替换为“AI豆”，且例外词保持原样，则输出“是”。<ret>- 如果存在至少一处非例外的“积分”未被替换，则输出“否”。<ret># 示例参考<ret>- 正例（输出“是”）：文本“您好！自2025年9月26日起，中国移动积分正式更名为AI豆。您当前有100AI豆，可前往积分商城兑换商品。”<ret>- 反例（输出“否”）：文本“您有100AI豆，积分可用于兑换话费。”<ret># 待校验文本：<ret>%s<ret><end><Bot>`
	prompt = fmt.Sprintf(prompt, r.content)
	client := openai.NewOpenAiClient(common.Model{
		Url:       "http://172.29.100.69:9969/v1/chat/completions",
		ApiKey:    "",
		ModelType: "x",
		IsReason:  false,
	}).
		SetStream(false).
		SetMaxTokens(4096).
		SetTemperature(0.6).
		SetTopK(0.99)
	response, err := client.Chat(prompt)
	if err != nil {
		r.resultData["备注"] = err.Error()
		return
	}
	answer := response.GetAnswer()
	r.resultData["rag_result"] = answer
}

func (r *RagTask) Do() {
	prompt := `<User>你是一名专业的文本校对AI助手，名为“意图匹配度检查”。对用户提供的意图名称与用户问进行快速、准确的校对。您以严谨、细致著称，只关注提供意图与用户问的匹配情况，不添加任何额外解释或情感色彩。您的唯一目标是输出“是”或“否”，确保结果客观可靠。<ret># 核心规则<ret>- 如果用户问属于所提供意图，输出【是】。<ret>- 如果用户问中明确包含意图名称文本，输出【是】。<ret>- 如果用户问不属于所提供意图或跟所提供意图完全不相干，输出【否】。<ret># 操作步骤<ret>1. 输入：接收用户提供的【意图名称】与【用户问】。<ret>2. 检查：<ret>- 分析提供的【用户问】是否属于用户提供的【意图】。<ret>3. 输出：<ret>- 如果符合，则输出【是】。<ret>- 如果不符合，则输出【否】。<ret># 意图名称：<ret>%s<ret># 用户问：<ret>%s<ret><end><Bot>`
	prompt = fmt.Sprintf(prompt, r.content, r.question)
	client := openai.NewOpenAiClient(common.Model{
		Url:       "http://172.29.100.69:9969/v1/chat/completions",
		ApiKey:    "",
		ModelType: "x",
		IsReason:  false,
	}).
		SetStream(false).
		SetMaxTokens(4096).
		SetTemperature(0.6).
		SetTopK(0.99)
	response, err := client.Chat(prompt)
	if err != nil {
		r.resultData["备注"] = err.Error()
		return
	}
	answer := response.GetAnswer()
	r.resultData["rag_result"] = answer
}

func (r *RagTask) Result() map[string]any {
	return maps.Clone(r.resultData)
}

type RagStress struct {
	question    string
	content     string
	inputLength int
	rounds      int
	resultData  map[string]any
}

func (r *RagStress) Reset() {
	r.inputLength = 0
	r.question = ""
	r.content = ""
	r.rounds = 0
	clear(r.resultData)
}

func (r *RagStress) NewTaskData(row map[string]any) error {
	// 复制数据
	r.resultData = maps.Clone(row)
	filedValues, err := tools.FindStringFromMapAny(r.resultData, "问题", "入参长度", "轮数", "知识ID")
	if err != nil {
		return err
	}
	r.inputLength, err = strconv.Atoi(filedValues[1])
	if err != nil {
		return err
	}
	r.rounds, err = strconv.Atoi(filedValues[2])
	if err != nil {
		return err
	}
	r.question = filedValues[0]
	if filedValues[3] != "" {
		bytes, err := os.ReadFile(filepath.Join("cache", filedValues[3]+".txt"))
		if err != nil {
			return err
		}
		r.content = string(bytes)
	} else {
		if cache == nil {
			mu.Lock()
			defer mu.Unlock()
			if cache == nil {
				cache, err = newKnowledgeTxt(filepath.Join("cache", "knowledge.txt"))
				if err != nil {
					return err
				}
			}
		}
		r.content = cache.getTxtByLength(r.inputLength)
	}

	return nil
}

func (r *RagStress) Do() {
	// 智能体信息
	agent := tools.Config.Agent["14B"]

	var children []map[string]any
	for i := 0; i < r.rounds; i++ {
		item := make(map[string]any)
		children = append(children, item)

		prompt := agent.DefaultSystemPrompt
		prompt = fmt.Sprintf(prompt, r.content, r.question)
		client := openai.NewOpenAiClient(common.Model{
			Url:       agent.Address,
			ApiKey:    agent.ApiKey,
			ModelType: agent.ModeType,
			IsReason:  false,
		}).
			SetStream(agent.Stream).
			SetMaxTokens(agent.MaxTokens).
			SetTemperature(agent.Temperature).
			SetTopK(agent.TopK)

		response, err := client.Chat(prompt)
		if err != nil {
			item["备注"] = err.Error()
			return
		}

		answer := response.GetAnswer()
		// 正则表达式：匹配 # Response 和 <end> 之间的内容
		// 使用 [\s\S] 或 [\d\D] 来匹配包括换行在内的所有字符
		re := regexp.MustCompile(`Response\s*([\s\S]*?)<end>`)

		// 查找匹配的内容
		matches := re.FindStringSubmatch(answer)

		if len(matches) >= 2 {
			answer = strings.TrimSpace(matches[1])
		}
		for _, rs := range []string{"# Type", "[Reply]", "# Response", "<end>", "[Reject]", "<ret>"} {
			answer = strings.ReplaceAll(answer, rs, "")
		}
		item["rag_result"] = strings.TrimSpace(answer)

		// 统计时间
		timeMap := response.GetTime()
		item["rag_time"] = timeMap["end"].Sub(timeMap["start"]).Milliseconds()
		item["rag_first_time"] = data.TernaryExpression(timeMap["first"].IsZero(), -1, timeMap["first"].Sub(timeMap["start"]).Milliseconds())

	}
	r.resultData["child"] = children
}

func (r *RagStress) Result() map[string]any {
	return maps.Clone(r.resultData)
}

var cache *knowledgeTxt
var mu sync.Mutex

type knowledgeTxt struct {
	content string
	mut     sync.Mutex
}

func newKnowledgeTxt(file string) (*knowledgeTxt, error) {
	k := &knowledgeTxt{mut: sync.Mutex{}}
	bytes, err := os.ReadFile(file)
	if err != nil {
		return nil, err
	}
	k.content = string(bytes)
	return k, nil
}

func (k *knowledgeTxt) getTxtByLength(length int) string {
	if length <= 0 {
		if k.content != "" {
			return k.content
		}
		return ""
	}
	if length > len(k.content) {
		n := length / len(k.content)
		a := length % len(k.content)
		return strings.Repeat(k.content, n) + k.content[:a]
	}
	return k.content[:length]
}
